# Name: Xueyan Lyu
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import datetime
from scipy.optimize import curve_fit

# read csv file
pr = pd.read_csv("owid-covid-data.csv")

# obtain data we need
UK_data = []
start_date = datetime.datetime.strptime("2021-05-25", '%Y-%m-%d')
end_date = datetime.datetime.strptime("2021-06-17", '%Y-%m-%d')
for i in range(len(pr)):
    line = pr.iloc[i]
    # UK cases
    if line['location'] == 'United Kingdom':
        date = line['date']
        # from 5.25 to 6.17
        if start_date <= datetime.datetime.strptime(date,'%Y-%m-%d') <= end_date:
            UK_data.append(line['new_cases'])


# define the exponential regression model
def fund(x, a, b):
    return a*(b**x)


# prepare the index
index = []
idate = start_date
index.append(str(idate.strftime('%m/%d')))
while idate < end_date:
    idate += datetime.timedelta(days=1)
    index.append(str(idate.strftime('%m/%d')))

# prepare data to draw the original plot
xdata = np.linspace(0, len(UK_data)-1, len(UK_data))
ydata = UK_data
# draw
plt.plot(index,ydata,'b-')
# fit the plot
# a = popt[0] b = popt[1]
popt, pcov = curve_fit(fund, xdata, ydata)
# prepare data to draw the fitting plot
ydata2 = [fund(i, popt[0],popt[1]) for i in xdata]
# draw
plt.plot(index,ydata2,'r--')
# print a and b
# the result is [2.88064241e+03 1.05450632e+00]
print(popt)
# show the plot
plt.xticks(rotation=60)  # rotate labels for x-axis
plt.show()